# -*- coding: utf-8 -*-
"""
A题最小可用解：导弹-云团遮蔽优化（LoS遮挡口径）
依赖：numpy（pip install numpy）

功能：
1) 物理/几何内核：导弹/无人机/云团位置、LoS-球体遮挡判定
2) 遮蔽时间段：扫描 + 二分法精修边界
3) 单团优化：粗网格(航向φ、速度v、引信τ) + 黄金分割搜起爆时刻t_det
4) 三团排程：候选池 + 贪心选3个，支持“同机释放间隔≥1s”

作者建议：
- “遮蔽”采用 LoS 口径：时刻t，导弹位置M(t)到目标T的线段与任一云团中心C_i(t)距离 ≤ r
- 默认参数：导弹300m/s；云团半径10m；云团下沉3m/s；有效20s

"""

from __future__ import annotations
import math
from dataclasses import dataclass, asdict
from typing import Callable, List, Tuple, Optional
import numpy as np

from tqdm import tqdm
# --------------------------
# 工具与基础几何
# --------------------------

Vec = np.ndarray

def unit(v) -> Vec:
    v = np.array(v, dtype=float)
    n = np.linalg.norm(v)
    return v if n == 0 else v / n

def closest_point_on_segment(A: Vec, B: Vec, C: Vec) -> Tuple[Vec, float]:
    """返回线段AB上距点C最近的点P与参数s∈[0,1]"""
    AB = B - A
    denom = float(np.dot(AB, AB))
    if denom == 0.0:
        return A, 0.0
    s = float(np.dot(C - A, AB) / denom)
    s = 0.0 if s < 0.0 else (1.0 if s > 1.0 else s)
    P = A + s * AB
    return P, s

# --------------------------
# 运动学模型
# --------------------------

@dataclass
class Missile:
    M0: Tuple[float, float, float]          # 初始位置
    speed: float = 300.0                    # m/s
    # 若不给方向，则默认朝向目标假点 fake_target
    direction: Optional[Tuple[float, float, float]] = None
    fake_target: Tuple[float, float, float] = (0.0, 0.0, 0.0)

    def pos(self, t: float, aim_to: Optional[Tuple[float,float,float]] = None) -> Vec:
        if self.direction is None:
            tgt = np.array(self.fake_target if aim_to is None else aim_to, float)
            d = unit(tgt - np.array(self.M0, float))
        else:
            d = unit(self.direction)
        return np.array(self.M0, float) + self.speed * d * t

@dataclass
class UAV:
    R0_xy: Tuple[float, float]              # 初始水平位置(x,y)
    z: float                                # 恒定飞行高度
    speed: float                            # m/s ∈ [70,140]
    heading_deg: float                      # 航向角，度（x正轴为0°，逆时针为正）

    def xy(self, t: float) -> Vec:
        phi = math.radians(self.heading_deg)
        u = np.array([math.cos(phi), math.sin(phi)])
        return np.array(self.R0_xy, float) + self.speed * u * t

@dataclass
class CloudModel:
    r: float = 10.0             # 云团半径(m)
    sink: float = 3.0           # 下沉速度(m/s)
    life: float = 20.0          # 有效时长(s)
    g: float = 9.81             # 重力加速度(用于投放下坠高度)

@dataclass
class DropPlan:
    """一枚干扰弹方案（绑定某架UAV）"""
    uav: UAV
    t_rel: float                # 释放时刻
    tau: float                  # 引信延时
    t_det: Optional[float] = None
    xy_det: Optional[Tuple[float,float]] = None
    z_det: Optional[float] = None

    def compute_detonation(self, model: CloudModel) -> None:
        self.t_det = self.t_rel + self.tau
        self.xy_det = tuple(self.uav.xy(self.t_det))
        self.z_det = self.uav.z - 0.5 * model.g * (self.tau ** 2)

    def center_func(self, model: CloudModel) -> Callable[[float], Optional[Vec]]:
        assert self.t_det is not None and self.xy_det is not None and self.z_det is not None
        def _f(t: float) -> Optional[Vec]:
            if t < self.t_det or t > self.t_det + model.life:
                return None
            x, y = self.xy_det
            z = self.z_det - model.sink * (t - self.t_det)
            return np.array([x, y, z], float)
        return _f

# --------------------------
# 遮蔽判定与时间段求解
# --------------------------

def los_distance(
    t: float,
    missile: Missile,
    target: Tuple[float,float,float],
    center_func: Callable[[float], Optional[Vec]]
) -> float:
    """时刻t，云团中心到线段 M(t)→T 的最小距离；若该时刻无云团则返回+inf"""
    C = center_func(t)
    if C is None:
        return float('inf')
    A = missile.pos(t)
    B = np.array(target, float)
    P, _ = closest_point_on_segment(A, B, C)
    return float(np.linalg.norm(P - C))

def find_cover_intervals(
    missile: Missile,
    target: Tuple[float,float,float],
    center_func: Callable[[float], Optional[Vec]],
    r: float,
    t0: float,
    t1: float,
    dt: float = 0.02,
    refine: bool = True,
    max_refine_iter: int = 50
) -> List[Tuple[float,float]]:
    """在[t0,t1]内找出 LoS 与半径r球体相交的时间段（覆盖区间）
       步骤：等步长扫描→符号变化处用二分精修边界
    """
    def f(t: float) -> float:
        return los_distance(t, missile, target, center_func) - r

    times, vals = [], []
    t = t0
    while t <= t1 + 1e-12:
        times.append(t)
        vals.append(f(t))
        t += dt

    intervals = []
    inside = vals[0] <= 0.0
    start_t = times[0] if inside else None

    def bisect_edge(a: float, b: float) -> float:
        fa, fb = f(a), f(b)
        for _ in range(max_refine_iter):
            m = 0.5 * (a + b)
            fm = f(m)
            if fa == 0.0:
                return a
            if fb == 0.0:
                return b
            if fa * fm <= 0.0:
                b, fb = m, fm
            else:
                a, fa = m, fm
        return 0.5 * (a + b)

    for i in range(1, len(times)):
        prev, cur = vals[i-1], vals[i]
        if inside and cur > 0.0:
            edge = bisect_edge(times[i-1], times[i]) if refine else times[i-1]
            intervals.append((start_t, edge))
            inside, start_t = False, None
        elif (not inside) and cur <= 0.0:
            start_t = bisect_edge(times[i-1], times[i]) if refine else times[i]
            inside = True

    if inside and start_t is not None:
        intervals.append((start_t, t1))

    # 合并接触/重叠
    merged = []
    for a,b in intervals:
        if not merged or a > merged[-1][1] + 1e-6:
            merged.append([a,b])
        else:
            merged[-1][1] = max(merged[-1][1], b)
    return [(a,b) for a,b in merged]

def union_length(intervals: List[Tuple[float,float]]) -> float:
    if not intervals:
        return 0.0
    ints = sorted(intervals)
    total = 0.0
    a0, b0 = ints[0]
    for a,b in ints[1:]:
        if a <= b0:
            b0 = max(b0, b)
        else:
            total += b0 - a0
            a0, b0 = a, b
    total += b0 - a0
    return total

# --------------------------
# 单弹优化：粗网格 + 黄金分割搜 t_det
# --------------------------

@dataclass
class SingleSearchConfig:
    phi_step_deg: float = 2.0                          # 航向角步长
    v_list: Tuple[float,...] = (70.0, 90.0, 110.0, 140.0)
    tau_list: Tuple[float,...] = (1.5, 2.5, 3.5, 4.5)  # 可改为连续搜索
    dt_scan: float = 0.02
    t_margin: float = 0.5                               # 边界余量

def golden_max(f: Callable[[float], float], a: float, b: float, iters: int = 40) -> Tuple[float,float]:
    """在[a,b]做黄金分割一维最大化，返回(x*, f(x*))"""
    gr = (math.sqrt(5) - 1) / 2
    c = b - gr * (b - a)
    d = a + gr * (b - a)
    fc, fd = f(c), f(d)
    for _ in range(iters):
        if fc > fd:
            b, d, fd = d, c, fc
            c = b - gr * (b - a)
            fc = f(c)
        else:
            a, c, fc = c, d, fd
            d = a + gr * (b - a)
            fd = f(d)
    x = (a + b) * 0.5
    return x, f(x)

def evaluate_single_plan(
    missile: Missile,
    target: Tuple[float,float,float],
    model: CloudModel,
    plan: DropPlan
) -> Tuple[float, List[Tuple[float,float]]]:
    """返回：遮蔽总时长、对应时间段列表"""
    plan.compute_detonation(model)
    cf = plan.center_func(model)
    t0 = plan.t_det
    t1 = plan.t_det + model.life
    intervals = find_cover_intervals(missile, target, cf, model.r, t0, t1, dt=0.02)
    return union_length(intervals), intervals

def optimize_single_cloud(
    missile: Missile,
    target: Tuple[float,float,float],
    uav: UAV,
    model: CloudModel,
    t_hit_max: Optional[float] = None,
    cfg: SingleSearchConfig = SingleSearchConfig()
) -> Tuple[DropPlan, float, List[Tuple[float,float]]]:
    """
    给定一架UAV，搜索单枚云团的最优方案：
    搜索变量：(heading=φ, speed=v, fuse=τ, detonation time=t_det[0, t_hit])
    约束：t_rel = t_det - τ ≥ 0
    """
    # 决定搜索时窗：导弹从M0到T的大致命中时间
    if t_hit_max is None:
        dist = float(np.linalg.norm(np.array(target, float) - np.array(missile.M0, float)))
        t_hit_max = dist / missile.speed

    best = (None, -1.0, [])  # (plan, score, intervals)

    phi_vals = np.arange(0.0, 360.0, cfg.phi_step_deg)
    for phi in phi_vals:
        for v in cfg.v_list:
            for tau in cfg.tau_list:
                # 黄金分割搜 t_det
                def score_at(t_det: float) -> float:
                    t_rel = t_det - tau
                    if t_rel < 0.0:  # 释放时刻不能为负
                        return -1e9
                    plan = DropPlan(uav=UAV(uav.R0_xy, uav.z, v, phi), t_rel=t_rel, tau=tau)
                    plan.compute_detonation(model)
                    cf = plan.center_func(model)
                    t0 = plan.t_det
                    t1 = min(plan.t_det + model.life, t_hit_max + cfg.t_margin)  # 过了命中后没意义
                    if t1 <= t0:
                        return -1e9
                    intervals = find_cover_intervals(missile, target, cf, model.r, t0, t1, dt=cfg.dt_scan)
                    return union_length(intervals)

                # 搜索区间：从0到t_hit_max
                x, fval = golden_max(score_at, 0.0, t_hit_max, iters=36)
                if fval > best[1]:
                    # 记录最优
                    t_det_star = x
                    t_rel_star = t_det_star - tau
                    plan_star = DropPlan(uav=UAV(uav.R0_xy, uav.z, v, phi), t_rel=t_rel_star, tau=tau)
                    score, intervals = evaluate_single_plan(missile, target, model, plan_star)
                    if score > best[1]:
                        best = (plan_star, score, intervals)

    assert best[0] is not None, "未找到可行方案，请扩大搜索范围。"
    return best  # plan, score, intervals

# --------------------------
# 三弹排程：候选池 + 贪心
# --------------------------

@dataclass
class GreedyConfig:
    topK_per_setting: int = 10
    phi_step_deg: float = 5.0
    v_list: Tuple[float,...] = (70.0, 100.0, 140.0)
    tau_list: Tuple[float,...] = (1.5, 2.5, 3.5, 4.5)
    t_grid_step: float = 0.2
    dt_scan: float = 0.02
    min_release_gap: float = 1.0   # 同机释放间隔约束（秒）

def marginal_gain(
    missile: Missile,
    target: Tuple[float,float,float],
    model: CloudModel,
    chosen: List[DropPlan],
    candidate: DropPlan,
    dt_scan: float = 0.02
) -> Tuple[float, List[Tuple[float,float]]]:
    """计算在已选方案 chosen 基础上，加入 candidate 的边际增益（并集长度增量）"""
    # 已选并集
    intervals_all = []
    for p in chosen:
        p.compute_detonation(model)
        cf = p.center_func(model)
        ints = find_cover_intervals(missile, target, cf, model.r, p.t_det, p.t_det + model.life, dt=dt_scan)
        intervals_all.extend(ints)
    base = union_length(intervals_all)

    # 加入候选
    candidate.compute_detonation(model)
    cf = candidate.center_func(model)
    ints_c = find_cover_intervals(missile, target, cf, model.r, candidate.t_det, candidate.t_det + model.life, dt=dt_scan)
    intervals_all.extend(ints_c)
    new_len = union_length(intervals_all)
    return new_len - base, ints_c

def schedule_three(
    missile: Missile,
    target: Tuple[float,float,float],
    uav: UAV,
    model: CloudModel,
    t_hit_max: Optional[float] = None,
    cfg: GreedyConfig = GreedyConfig()
) -> Tuple[List[DropPlan], float, List[Tuple[float,float]]]:
    """同一架UAV 投放三弹，贪心最大化遮蔽并集；含释放≥1s约束"""
    if t_hit_max is None:
        dist = float(np.linalg.norm(np.array(target, float) - np.array(missile.M0, float)))
        t_hit_max = dist / missile.speed

    # 生成候选池（粗网格）
    candidates: List[DropPlan] = []
    for phi in np.arange(0.0, 360.0, cfg.phi_step_deg):
        for v in cfg.v_list:
            for tau in cfg.tau_list:
                scores_local = []
                for t_det in np.arange(0.0, t_hit_max, cfg.t_grid_step):
                    t_rel = t_det - tau
                    if t_rel < 0.0:
                        continue
                    plan = DropPlan(uav=UAV(uav.R0_xy, uav.z, v, phi), t_rel=t_rel, tau=tau)
                    # 快速打分（粗）：用扫描步长dt_scan
                    s, _ = evaluate_single_plan(missile, target, model, plan)
                    if s > 0:
                        scores_local.append((s, plan))
                # 每组(phi,v,tau)保留 topK
                scores_local.sort(key=lambda x: x[0], reverse=True)
                for s, p in scores_local[:cfg.topK_per_setting]:
                    candidates.append(p)

    # 贪心选3个
    chosen: List[DropPlan] = []
    chosen_intervals: List[Tuple[float,float]] = []
    total_len = 0.0

    while len(chosen) < 3 and candidates:
        best_gain, best_p, best_ints = -1.0, None, []
        for c in candidates:
            # 释放间隔约束（同机）
            ok = True
            for p in chosen:
                if abs((c.t_rel) - (p.t_rel)) < cfg.min_release_gap:
                    ok = False
                    break
            if not ok:
                continue
            gain, ints = marginal_gain(missile, target, model, chosen, c, dt_scan=cfg.dt_scan)
            if gain > best_gain:
                best_gain, best_p, best_ints = gain, c, ints

        if best_p is None or best_gain <= 1e-6:
            break
        chosen.append(best_p)
        # 更新总并集长度
        chosen_intervals.extend(best_ints)
        total_len = union_length(chosen_intervals)
        # 删除与已选释放时间太近的候选
        candidates = [q for q in candidates if all(abs(q.t_rel - p.t_rel) >= cfg.min_release_gap for p in chosen)]

    return chosen, total_len, chosen_intervals

# --------------------------
# 示例：题1固定参数 / 题2单团优化 / 三团贪心
# --------------------------

def demo():
    # ===== 场景参数（可按题面修改） =====
    T = (0.0, 200.0, 5.0)                       # 目标代表点
    missile = Missile(M0=(1500.0, -1200.0, 800.0), speed=300.0, direction=None, fake_target=(0.0,0.0,0.0))
    uav = UAV(R0_xy=(-600.0, -1500.0), z=300.0, speed=120.0, heading_deg=35.0)
    model = CloudModel(r=10.0, sink=3.0, life=20.0, g=9.81)

    # ===== 题1：给定 v=120, t_rel=1.5s, tau=3.6s, heading=35° =====
    plan1 = DropPlan(uav=uav, t_rel=1.5, tau=3.6)
    length1, intervals1 = evaluate_single_plan(missile, T, model, plan1)
    print("\n[题1固定参] 遮蔽总时长(s):", round(length1, 3), " 时间段:", [(round(a,3), round(b,3)) for a,b in intervals1])

    # ===== 题2：单团优化 =====
    best_plan, best_len, best_ints = optimize_single_cloud(missile, T, uav, model)
    print("\n[题2单团最优] 遮蔽总时长(s):", round(best_len, 3))
    print("  方案:", {**asdict(best_plan), "heading_deg": best_plan.uav.heading_deg, "uav_speed": best_plan.uav.speed})
    print("  时间段:", [(round(a,3), round(b,3)) for a,b in best_ints])

    # ===== 题3：三团贪心排程 =====
    chosen, total_len, chosen_ints = schedule_three(missile, T, uav, model)
    print("\n[题3三团贪心] 遮蔽总时长(s):", round(total_len, 3))
    for i,p in enumerate(chosen,1):
        print(f"  弹{i}: t_rel={p.t_rel:.2f}s, tau={p.tau:.2f}s, heading={p.uav.heading_deg:.1f}°, speed={p.uav.speed:.1f}m/s")
    print("  并集时间段:", [(round(a,3), round(b,3)) for a,b in chosen_ints])

if __name__ == "__main__":
    demo()
